package solvers.linearsupport;

import globals.ExperimentTracker;
import globals.StopWatch;

import java.util.ArrayList;

import problems.DiscreteActions;
import problems.DiscreteObservations;
import problems.Mopomdp;
import solvers.vi.BeliefSampler;
import solvers.vi.ValueIteratorPO;
import value.*;


public class VanillaPerseusSolver<S, A, O> implements SingleSolver<A> {

	public Mopomdp<S,A,O> problem;
	private int nSamples;
	private int maxIterations; 
	private double precision;
	public ArrayList<double[]> bSampleSet;
	private double[] initialBelief; 
	private int cnt=0;
	
	private static final boolean LOG = true;
	private static final boolean DEBUG = true;
	
	public VanillaPerseusSolver(Mopomdp<S,A,O> m, double[] iBelief, int nSamplez, int maxIter, double preciz){
		this.problem = m;
		this.bSampleSet = null;
		this.maxIterations = maxIter;
		this.nSamples = nSamplez;
		this.precision = preciz;
		this.initialBelief = iBelief;
		this.cnt = 0;
	}
	
	public VanillaPerseusSolver(Mopomdp<S,A,O> m, double[] iBelief, ArrayList<double[]>  bs, int maxIter, double preciz){
		this.problem = m;
		this.bSampleSet = bs;
		this.maxIterations = maxIter;
		this.nSamples = bs.size();
		this.precision = preciz;
		this.initialBelief = iBelief;
		this.cnt = 0;
	}
	
	public void initSamples(){
		if(bSampleSet == null){
			BeliefSampler<S,A,O> bs = new BeliefSampler<S,A,O>(this.problem);
			this.bSampleSet = bs.randomExplorationPerseus(nSamples);
		}
	}
	
	
	
	@Override
	public ValueAtBeliefVector<A> solveForW(CornerPoint<A> cp) {
		return this.solveForW(cp.weightVector);
	}

	@SuppressWarnings("unchecked")
	@Override
	public ValueAtBeliefVector<A> solveForW(double[] w) {
		
		if(DEBUG){
			System.out.print("Solver called with Weight: [");
			for(int i=0; i<w.length; i++){
				System.out.print(w[i]+"  ");
			}
			System.out.println("]");
		}
		
		this.initSamples();
		
		ArrayList<LazyScalarizedVector<A>> lb = new ArrayList<LazyScalarizedVector<A>>();
		LazyScalarizedVector<A> lsv = new LazyScalarizedVector<A>(this.problem.misereMatrix(), w);
		lb.add(lsv);
		ValueIteratorPO<LazyScalarizedVector<A>,S,A,O> vi = 
				new ValueIteratorPO<LazyScalarizedVector<A>,S,A,O>(this.problem, lb);
		try{
			DiscreteActions<A,S> dacts   = (DiscreteActions<A,S>) problem.getActionSet();
			DiscreteObservations<O> dobs = (DiscreteObservations<O>) problem.getObservationSet();
			vi.perseus(bSampleSet, dacts, dobs, problem.getStateSet(), this.precision, this.maxIterations);
			ValueAtBeliefVector<A> vb = new ValueAtBeliefVector<A>(w.length, this.initialBelief, vi.alphas);
			if(LOG){
				StopWatch sw = StopWatch.getGlobalStopWatch();
				sw.stop();
				ExperimentTracker et = ExperimentTracker.getGlobalTracker(null);
				ArrayList<ValueAtBeliefVector<A>> vbs = new ArrayList<ValueAtBeliefVector<A>>(1);
				vbs.add(vb);
				et.writeVectorSet(""+(++cnt)+"_", vbs);
				System.out.println();
				et.trackNewVec(vb);
				Thread.sleep(100);
				sw.start();
			}
			return vb;
		} catch(Exception e) {
			System.err.println(e.toString());
			System.err.println("Trying to call discrete POMDP method for non-discrete problem.");
			System.err.println("Quitting...");
			System.exit(77);
			return null;
		}
		
	}

}
